import torch
from torch.utils.data import Dataset
import numpy as np
from skimage import io as skio
from skimage.transform import resize, rotate
import os


def default_loader(path):
    return resize(skio.imread(path), (512, 512, 3))


class MyDataset(Dataset):
    def __init__(self, img_file_list, mask_file_list, img_file_dir, mask_file_dir, loader=default_loader, transform=False):
        self.image_file_list = img_file_list
        self.mask_file_list = mask_file_list
        self.loader = loader
        self.img_file_dir = img_file_dir
        self.mask_file_dir = mask_file_dir
        self.transform = transform

    def __getitem__(self, index):
        path1 = self.image_file_list[index]
        img_path = os.path.join(self.img_file_dir, path1)
        img = self.loader(img_path)
        path2 = self.mask_file_list[index]
        mask_path = os.path.join(self.mask_file_dir, path2)
        mask = self.loader(mask_path)
        mask = (mask > 0).astype(np.uint8)
        if self.transform:
            degree = np.random.randint(0, 360)
            img = rotate(img, degree, resize=False)
            mask = rotate(img, degree, resize=False)
        # return img, mask
        return {
            'img': img,
            'mask': mask
        }

    def __len__(self):
        return len(self.image_file_list)




